from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import declarative_base
from sqlalchemy import select
from typing import AsyncGenerator
import os

# 尝试从环境变量获取数据库URL，否则使用默认SQLite
DATABASE_URL = os.getenv("DATABASE_URL", "sqlite+aiosqlite:///./stock_data.db")

# 创建异步引擎
# echo=True 会在控制台输出SQL语句，便于调试
engine = create_async_engine(
    DATABASE_URL,
    echo=True,
)

# 创建异步会话工厂
# expire_on_commit=False 避免在提交后session过期
AsyncSessionLocal = async_sessionmaker(
    engine, expire_on_commit=False
)

# 创建基础模型类
Base = declarative_base()


# 获取数据库会话的依赖项
async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """
    数据库会话依赖项，用于FastAPI的Depends
    
    Yields:
        AsyncSession: 异步数据库会话对象
    """
    async with AsyncSessionLocal() as session:
        yield session

# 初始化数据库的异步函数
async def init_db():
    """
    初始化数据库，创建所有表
    """
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    print("✅ 数据库表已成功创建")

# 关闭数据库连接
async def close_db():
    """关闭数据库连接"""
    await engine.dispose()